from boson_multimodal.serve.serve_engine import HiggsAudioServeEngine, HiggsAudioResponse
from boson_multimodal.data_types import ChatMLSample, Message, AudioContent

import torch
import torchaudio
import time
import click

MODEL_PATH = "bosonai/higgs-audio-v2-generation-3B-base"
AUDIO_TOKENIZER_PATH = "bosonai/higgs-audio-v2-tokenizer"

system_prompt = (
    "Generate audio following instruction.\n\n<|scene_desc_start|>\nAudio is recorded from a quiet room.\n<|scene_desc_end|>"
)

messages = [
    Message(
        role="system",
        content=system_prompt,
    ),
    Message(
        role="user",
        content="我个人认为意大利面就应该拌四十二号混泥土，因为这个螺丝钉的长度他是很容易直接影响到挖掘机的扭矩，你知道吧，你往里砸的时候，一瞬间就会产生大量的高能蛋白，俗称UFO，会严重影响经济的发展，甚至对一个太平洋以及充电器造成一定的核污染，你知道吧，啊？再者说根据这个勾股定理，你可以很容易的推断出，人工饲养的东条鹰鸡，他是可以捕获野生的三角函数的，所以这个不管秦始皇的切面是否具有放射性，特朗普的n次方是否具有沉淀物，都不影响这个沃尔玛跟这个维尔康在南极汇合",
    ),
]
device = "cuda" if torch.cuda.is_available() else "cpu"

serve_engine = HiggsAudioServeEngine(MODEL_PATH, AUDIO_TOKENIZER_PATH, device=device)

output: HiggsAudioResponse = serve_engine.generate(
    chat_ml_sample=ChatMLSample(messages=messages),
    max_new_tokens=1024,
    temperature=0.3,
    top_p=0.95,
    top_k=50,
    stop_strings=["<|end_of_text|>", "<|eot_id|>"],
)
torchaudio.save(f"chinese_output.wav", torch.from_numpy(output.audio)[None, :], output.sampling_rate)
